package model

import (
	"cluster-cache/common"
	"github.com/cespare/xxhash/v2"
	tr "github.com/emirpasic/gods/trees/redblacktree"
	"github.com/emirpasic/gods/utils"
	"strconv"
	"sync"
)

type SortSet interface {
	Add(key string, score float64, v string) error
	Delete(key string, v string) error
	Range(key string, min string, max string) ([]string, []string)
}

type MultiSortSet struct {
	store []map[string]*TreeSet
	locks []*sync.RWMutex
	size  uint64
}

func NewMultiSortSet(size int) SortSet {
	ms := &MultiSortSet{
		store: make([]map[string]*TreeSet, size),
		locks: make([]*sync.RWMutex, size),
		size:  uint64(size),
	}
	for i := 0; i < size; i++ {
		ms.store[i] = make(map[string]*TreeSet)
		ms.locks[i] = new(sync.RWMutex)
	}
	return ms
}

func (m *MultiSortSet) Add(key string, score float64, v string) error {
	bucket := xxhash.Sum64(common.StringToBytes(key)) % m.size
	m.locks[bucket].Lock()
	if val, ok := m.store[bucket][key]; ok {
		val.Add(v, int64(score))
	} else {
		val = NewTreeSet()
		val.Add(v, int64(score))
		m.store[bucket][key] = val
	}
	m.locks[bucket].Unlock()
	return nil
}

func (m *MultiSortSet) Delete(key string, v string) error {
	bucket := xxhash.Sum64(common.StringToBytes(key)) % m.size
	m.locks[bucket].Lock()
	if val, ok := m.store[bucket][key]; ok {
		val.Delete(v)
	}
	m.locks[bucket].Unlock()
	return nil
}

func (m *MultiSortSet) Range(key string, min string, max string) ([]string, []string) {
	bucket := xxhash.Sum64(common.StringToBytes(key)) % m.size
	m.locks[bucket].RLock()
	if val, ok := m.store[bucket][key]; ok {
		minScore, err := strconv.ParseFloat(min, 64)
		if err != nil {
			m.locks[bucket].RUnlock()
			return nil, nil
		}
		maxScore, err := strconv.ParseFloat(max, 64)
		if err != nil {
			m.locks[bucket].RUnlock()
			return nil, nil
		}
		m.locks[bucket].RUnlock()
		return val.Range(int64(minScore), int64(maxScore))
	} else {
		m.locks[bucket].RUnlock()
		return nil, nil
	}
}

type TreeSet struct {
	sync.RWMutex
	tree        *tr.Tree
	value2score map[string]int64
}

func NewTreeSet() *TreeSet {
	return &TreeSet{
		tree:        tr.NewWith(utils.Int64Comparator),
		value2score: make(map[string]int64),
	}
}

func (t *TreeSet) Add(v string, score int64) {
	t.Lock()
	if old, ok := t.value2score[v]; ok {
		if old == score {
			t.Unlock()
			return
		}
		t.tree.Remove(old)
		t.tree.Put(score, v)
		t.value2score[v] = score
	} else {
		t.tree.Put(score, v)
		t.value2score[v] = score
	}
	t.Unlock()
}

func (t *TreeSet) Delete(v string) {
	t.Lock()
	if old, ok := t.value2score[v]; ok {
		t.tree.Remove(old)
		delete(t.value2score, v)
	}
	t.Unlock()
}

func (t *TreeSet) Range(min int64, max int64) ([]string, []string) {
	t.RLock()
	node, ok := t.tree.Ceiling(min)
	values := make([]string, 0)
	scores := make([]string, 0)
	at := t.tree.IteratorAt(node)
	for ok {
		if at.Key().(int64) > max {
			break
		}
		scores = append(scores, strconv.FormatInt(at.Key().(int64), 10))
		values = append(values, at.Value().(string))
		ok = at.Next()
	}
	t.RUnlock()
	return values, scores
}
